# -*- coding = utf-8 -*-
# @Time : 2022/3/9 18:22
# @Author : GHHHHHHHHH
# @File : train.py
# @Software : PyCharm
import matplotlib.pyplot as plt
import torch
from tqdm import tqdm

from model import Model, MseLossPro


class TrainBiasFunkSVD:
    def __init__(self, user, item, dataset, hidden=None, epoch=50, LR=0.01):
        self.user = user
        self.item = item
        self.hidden = hidden
        self.epoch = epoch
        self.LR = LR
        self.dataset = dataset

    def train(self):
        if self.hidden is None:
            for k in tqdm(range(1, max(self.user, self.item) * 5)):
                self.__train(k)
        else:
            print("Error:", self.__train(self.hidden))

    def __train(self, k):
        # 损失函数
        criterion = MseLossPro(gamma=0.0075)
        # 步长
        LR = self.LR
        model = Model(self.user, self.item, k)
        best_loss = 100000
        losses = []
        for epoch in range(self.epoch):
            optimizer = torch.optim.Adam(model.parameters(), lr=LR)
            y_pred = model()
            loss = criterion(y_pred, torch.Tensor(self.dataset), model.user_hidden, model.hidden_item,
                             model.user_bias, model.item_bias)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            losses.append(loss.item())
            if loss.item() < best_loss:
                best_loss = loss.item()
                torch.save(model, "model.pth")
        plt.plot(losses)
        plt.show()
        return best_loss
